"""
对nii的标注数据进行2D切片，并保存为jpeg文件
"""
from lib.TransFuse import TransFuse_S
import os
import sys
from PIL import Image
sys.path.append(os.path.split(sys.path[0])[0])
import numpy as np
from config import zunYiParameter as para

from dataset.zunYi.dataloader import getTestLoader
import torch
def saveSeg2D(true_seg_array,pre_seg_array,save_path,seg_index):
    '''
    将array保存为图片
    seg_index是图片的序号，保存信息，方便后续工作
    '''
    true_seg_array*=255
    pre_seg_array*=255
    file_type='.png'
    true_seg_save_path=os.path.join(save_path,str(seg_index),'true'+file_type)
    pre_seg_save_path=os.path.join(save_path,str(seg_index),'pre'+file_type)
    if not os.path.exists(os.path.dirname(true_seg_save_path)):
        os.makedirs(os.path.dirname(true_seg_save_path))
    if not os.path.exists(os.path.dirname(pre_seg_save_path)):
        os.makedirs(os.path.dirname(pre_seg_save_path))
    true_seg_array=true_seg_array.astype(np.uint8)
    pre_seg_array = pre_seg_array.astype(np.uint8)

    true_seg_array=np.squeeze(true_seg_array)
    pre_seg_array = np.squeeze(pre_seg_array)


    im_true=Image.fromarray(true_seg_array)
    im_true=im_true.convert('L')
    im_pre=Image.fromarray(pre_seg_array)
    im_pre=im_pre.convert('L')

    im_true.save(true_seg_save_path)
    im_pre.save(pre_seg_save_path)

def evaluate(model):

    model.eval()

    test_loader = getTestLoader(para.cut2d_save_path,id_path = para.test2d_id_path,batchsize = 1)

    n_class=2

    for i,(image, gt) in enumerate(test_loader):
        image = image.cuda()

        with torch.no_grad():
            _, _, res = model(image)
        # print(image.shape,res.shape,torch.tensor(gt).unsqueeze(0).unsqueeze(0).shape)
        # loss = structure_loss(res, torch.tensor(gt).unsqueeze(0).unsqueeze(0).cuda())

        res = res.sigmoid().data.cpu().numpy().squeeze()
        gt = 1*(gt>0.5)
        res = 1*(res > 0.5)

        gt=gt.cpu().numpy()
        if (1 in gt) or (1 in res):
            saveSeg2D(gt,res,para.photo_save_path,i)
        print('save one photo')
    return None





if __name__ == '__main__':
    model_path = '/home/liukai/projects/TransFuseForBreast/snapshots/TransFuse_S/TransFuse-89.pth'
    model = TransFuse_S(in_chans = 1, img_size = 448).cuda()
    model.load_state_dict(torch.load(model_path))
    model.cuda()
    model.eval()
    result = evaluate(model)
    print(result)
